// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/util/align_util.h>

#include <nebula/core/array.h>
#include <nebula/core/chunked_array.h>
#include <nebula/core/record_batch.h>
#include <nebula/core/table.h>
#include <nebula/types/type_fwd.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula::util {
    bool CheckAlignment(const Buffer &buffer, int64_t alignment) {
        if (alignment <= 0) {
            return true;
        }
        return buffer.address() % alignment == 0;
    }

    namespace {

        // Returns the type that controls how the buffers of this ArrayData (not its children)
        // should behave
        Type::type GetTypeForBuffers(const ArrayData &array) {
            Type::type type_id = array.type->storage_id();
            if (type_id == Type::DICTIONARY) {
                return ::turbo::checked_pointer_cast<DictionaryType>(array.type)
                        ->index_type()
                        ->id();
            }
            return type_id;
        }

        // Checks to see if an array's own buffers are aligned but doesn't check
        // children
        bool CheckSelfAlignment(const ArrayData &array, int64_t alignment) {
            if (alignment == kValueAlignment) {
                Type::type type_id = GetTypeForBuffers(array);
                for (std::size_t i = 0; i < array.buffers.size(); i++) {
                    if (array.buffers[i]) {
                        int expected_alignment =
                                RequiredValueAlignmentForBuffer(type_id, static_cast<int>(i));
                        if (!CheckAlignment(*array.buffers[i], expected_alignment)) {
                            return false;
                        }
                    }
                }
            } else {
                for (const auto &buffer: array.buffers) {
                    if (buffer) {
                        if (!CheckAlignment(*buffer, alignment)) return false;
                    }
                }
            }
            return true;
        }

    }  // namespace

    bool CheckAlignment(const ArrayData &array, int64_t alignment) {
        if (!CheckSelfAlignment(array, alignment)) {
            return false;
        }

        if (array.dictionary) {
            if (!CheckAlignment(*array.dictionary, alignment)) return false;
        }

        for (const auto &child: array.child_data) {
            if (child) {
                if (!CheckAlignment(*child, alignment)) return false;
            }
        }
        return true;
    }

    bool CheckAlignment(const Array &array, int64_t alignment) {
        return CheckAlignment(*array.data(), alignment);
    }

    bool CheckAlignment(const ChunkedArray &array, int64_t alignment,
                        std::vector<bool> *needs_alignment, int offset) {
        bool all_aligned = true;
        needs_alignment->resize(needs_alignment->size() + array.num_chunks(), false);
        for (auto i = 0; i < array.num_chunks(); ++i) {
            if (array.chunk(i) && !CheckAlignment(*array.chunk(i), alignment)) {
                (*needs_alignment)[i + offset] = true;
                all_aligned = false;
            }
        }
        return all_aligned;
    }

    bool CheckAlignment(const RecordBatch &batch, int64_t alignment,
                        std::vector<bool> *needs_alignment) {
        bool all_aligned = true;
        needs_alignment->resize(batch.num_columns(), false);
        for (auto i = 0; i < batch.num_columns(); ++i) {
            if (batch.column(i) && !CheckAlignment(*batch.column(i), alignment)) {
                (*needs_alignment)[i] = true;
                all_aligned = false;
            }
        }
        return all_aligned;
    }

    bool CheckAlignment(const Table &table, int64_t alignment,
                        std::vector<bool> *needs_alignment) {
        bool all_aligned = true;
        needs_alignment->resize(table.num_columns(), false);
        for (auto i = 1; i <= table.num_columns(); ++i) {
            if (table.column(i - 1) &&
                !CheckAlignment(*table.column(i - 1), alignment, needs_alignment,
                                (i - 1) * (1 + table.column(i - 1)->num_chunks()))) {
                (*needs_alignment)[i * table.column(i - 1)->num_chunks() + i - 1] = true;
                all_aligned = false;
            }
        }
        return all_aligned;
    }

    turbo::Result<std::shared_ptr<Buffer>> EnsureAlignment(std::shared_ptr<Buffer> buffer,
                                                    int64_t alignment,
                                                    MemoryPool *memory_pool) {
        if (alignment == kValueAlignment) {
            return turbo::invalid_argument_error(
                    "The kValueAlignment option may only be used to call EnsureAlignment on arrays "
                    "or tables and cannot be used with buffers");
        }
        if (alignment <= 0) {
            return turbo::invalid_argument_error("Alignment must be a positive integer");
        }
        if (!CheckAlignment(*buffer, alignment)) {
            if (!buffer->is_cpu()) {
                return turbo::unimplemented_error("Reallocating an unaligned non-CPU buffer.");
            }
            int64_t minimum_desired_alignment = std::max(kDefaultBufferAlignment, alignment);
            TURBO_MOVE_OR_RAISE(
                    auto new_buffer,
                    allocate_buffer(buffer->size(), minimum_desired_alignment, memory_pool));
            std::memcpy(new_buffer->mutable_data(), buffer->data(), buffer->size());
            // R build with openSUSE155 requires an explicit shared_ptr construction
            return std::shared_ptr<Buffer>(std::move(new_buffer));
        } else {
            return buffer;
        }
    }

    turbo::Result<std::shared_ptr<ArrayData>> EnsureAlignment(std::shared_ptr<ArrayData> array_data,
                                                       int64_t alignment,
                                                       MemoryPool *memory_pool) {
        if (!CheckAlignment(*array_data, alignment)) {
            std::vector<std::shared_ptr<Buffer>> buffers = array_data->buffers;
            Type::type type_id = GetTypeForBuffers(*array_data);
            for (size_t i = 0; i < buffers.size(); ++i) {
                if (buffers[i]) {
                    int64_t expected_alignment = alignment;
                    if (alignment == kValueAlignment) {
                        expected_alignment =
                                RequiredValueAlignmentForBuffer(type_id, static_cast<int>(i));
                    }
                    TURBO_MOVE_OR_RAISE(
                            buffers[i],
                            EnsureAlignment(std::move(buffers[i]), expected_alignment, memory_pool));
                }
            }

            for (auto &it: array_data->child_data) {
                TURBO_MOVE_OR_RAISE(it, EnsureAlignment(std::move(it), alignment, memory_pool));
            }

            if (array_data->type->id() == Type::DICTIONARY) {
                TURBO_MOVE_OR_RAISE(
                        array_data->dictionary,
                        EnsureAlignment(std::move(array_data->dictionary), alignment, memory_pool));
            }

            auto new_array_data = ArrayData::create(
                    array_data->type, array_data->length, std::move(buffers), array_data->child_data,
                    array_data->dictionary, array_data->get_null_count(), array_data->offset);
            return new_array_data;
        } else {
            return array_data;
        }
    }

    turbo::Result<std::shared_ptr<Array>> EnsureAlignment(std::shared_ptr<Array> array,
                                                   int64_t alignment,
                                                   MemoryPool *memory_pool) {
        TURBO_MOVE_OR_RAISE(auto new_array_data,
                               EnsureAlignment(array->data(), alignment, memory_pool));

        if (new_array_data.get() == array->data().get()) {
            return array;
        } else {
            return make_array(std::move(new_array_data));
        }
    }

    turbo::Result<std::shared_ptr<ChunkedArray>> EnsureAlignment(std::shared_ptr<ChunkedArray> array,
                                                          int64_t alignment,
                                                          MemoryPool *memory_pool) {
        std::vector<bool> needs_alignment;
        if (!CheckAlignment(*array, alignment, &needs_alignment)) {
            ArrayVector chunks_ = array->chunks();
            for (int i = 0; i < array->num_chunks(); ++i) {
                if (needs_alignment[i] && chunks_[i]) {
                    TURBO_MOVE_OR_RAISE(
                            chunks_[i], EnsureAlignment(std::move(chunks_[i]), alignment, memory_pool));
                }
            }
            return ChunkedArray::create(std::move(chunks_), array->type());
        } else {
            return array;
        }
    }

    turbo::Result<std::shared_ptr<RecordBatch>> EnsureAlignment(std::shared_ptr<RecordBatch> batch,
                                                         int64_t alignment,
                                                         MemoryPool *memory_pool) {
        std::vector<bool> needs_alignment;
        if (!CheckAlignment(*batch, alignment, &needs_alignment)) {
            ArrayVector columns_ = batch->columns();
            for (int i = 0; i < batch->num_columns(); ++i) {
                if (needs_alignment[i] && columns_[i]) {
                    TURBO_MOVE_OR_RAISE(
                            columns_[i], EnsureAlignment(std::move(columns_[i]), alignment, memory_pool));
                }
            }
            return RecordBatch::create(batch->schema(), batch->num_rows(), std::move(columns_));
        } else {
            return batch;
        }
    }

    turbo::Result<std::shared_ptr<Table>> EnsureAlignment(std::shared_ptr<Table> table,
                                                   int64_t alignment,
                                                   MemoryPool *memory_pool) {
        std::vector<bool> needs_alignment;
        if (!CheckAlignment(*table, alignment, &needs_alignment)) {
            std::vector<std::shared_ptr<ChunkedArray>> columns_ = table->columns();
            for (int i = 1; i <= table->num_columns(); ++i) {
                if (columns_[i - 1] && needs_alignment[i * columns_[i - 1]->num_chunks() + i - 1]) {
                    ArrayVector chunks_ = columns_[i - 1]->chunks();
                    for (size_t j = 0; j < chunks_.size(); ++j) {
                        if (chunks_[j] &&
                            needs_alignment[j + (i - 1) * (1 + columns_[i - 1]->num_chunks())]) {
                            TURBO_MOVE_OR_RAISE(chunks_[j], EnsureAlignment(std::move(chunks_[j]),
                                                                               alignment, memory_pool));
                        }
                    }
                    TURBO_MOVE_OR_RAISE(
                            columns_[i - 1],
                            ChunkedArray::create(std::move(chunks_), columns_[i - 1]->type()));
                }
            }
            return Table::create(table->schema(), std::move(columns_), table->num_rows());
        } else {
            return table;
        }
    }

}  // namespace nebula::util
